from django.db import models
from django.conf import settings
import hashlib


class KnowledgeBase(models.Model):
    """知识库"""
    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, verbose_name='用户')
    name = models.CharField(max_length=100, verbose_name='知识库名称')
    description = models.TextField(null=True, blank=True, verbose_name='描述')
    icon = models.CharField(max_length=50, default='📚', verbose_name='图标')
    is_default = models.BooleanField(default=False, verbose_name='是否默认知识库')
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    updated_at = models.DateTimeField(auto_now=True, verbose_name='更新时间')

    class Meta:
        db_table = 'knowledge_base'
        verbose_name = '知识库'
        verbose_name_plural = verbose_name
        ordering = ['-created_at']
        unique_together = ['user', 'name']

    def __str__(self):
        return f'{self.user.username} - {self.name}'


class Tag(models.Model):
    """标签"""
    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, verbose_name='用户')
    name = models.CharField(max_length=50, verbose_name='标签名称')
    color = models.CharField(max_length=7, default='#409EFF', verbose_name='标签颜色')
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')

    class Meta:
        db_table = 'tag'
        verbose_name = '标签'
        verbose_name_plural = verbose_name
        ordering = ['-created_at']
        unique_together = ['user', 'name']

    def __str__(self):
        return self.name


class Document(models.Model):
    """文档"""
    FILE_TYPE_CHOICES = [
        ('pdf', 'PDF'),
        ('docx', 'Word文档'),
        ('txt', '文本文件'),
    ]

    STATUS_CHOICES = [
        ('uploading', '上传中'),
        ('processing', '处理中'),
        ('completed', '已完成'),
        ('failed', '失败'),
    ]

    knowledge_base = models.ForeignKey(KnowledgeBase, on_delete=models.CASCADE, verbose_name='知识库')
    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, verbose_name='用户')
    title = models.CharField(max_length=200, verbose_name='文档标题')
    file = models.FileField(upload_to='documents/%Y/%m/', verbose_name='文件')
    file_type = models.CharField(max_length=10, choices=FILE_TYPE_CHOICES, verbose_name='文件类型')
    file_size = models.BigIntegerField(verbose_name='文件大小(字节)')
    md5_hash = models.CharField(max_length=32, verbose_name='MD5哈希', db_index=True)
    page_count = models.IntegerField(null=True, blank=True, verbose_name='页数')
    word_count = models.IntegerField(null=True, blank=True, verbose_name='字数')
    status = models.CharField(max_length=20, choices=STATUS_CHOICES, default='uploading', verbose_name='状态')
    error_message = models.TextField(null=True, blank=True, verbose_name='错误信息')
    tags = models.ManyToManyField(Tag, blank=True, verbose_name='标签')
    content = models.TextField(null=True, blank=True, verbose_name='提取的文本内容')
    vector_indexed = models.BooleanField(default=False, verbose_name='是否已向量化')
    vector_id = models.CharField(max_length=100, null=True, blank=True, verbose_name='向量数据库ID')
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')
    updated_at = models.DateTimeField(auto_now=True, verbose_name='更新时间')

    class Meta:
        db_table = 'document'
        verbose_name = '文档'
        verbose_name_plural = verbose_name
        ordering = ['-created_at']
        indexes = [
            models.Index(fields=['user', 'status']),
            models.Index(fields=['knowledge_base', 'status']),
        ]

    def __str__(self):
        return self.title

    @staticmethod
    def calculate_md5(file_obj):
        """计算文件MD5"""
        md5 = hashlib.md5()
        for chunk in file_obj.chunks():
            md5.update(chunk)
        return md5.hexdigest()


class DocumentChunk(models.Model):
    """文档分块"""
    document = models.ForeignKey(Document, on_delete=models.CASCADE, related_name='chunks', verbose_name='文档')
    chunk_index = models.IntegerField(verbose_name='块索引')
    content = models.TextField(verbose_name='内容')
    token_count = models.IntegerField(verbose_name='Token数')
    vector_id = models.CharField(max_length=100, null=True, blank=True, verbose_name='向量ID')
    created_at = models.DateTimeField(auto_now_add=True, verbose_name='创建时间')

    class Meta:
        db_table = 'document_chunk'
        verbose_name = '文档分块'
        verbose_name_plural = verbose_name
        ordering = ['document', 'chunk_index']
        unique_together = ['document', 'chunk_index']

    def __str__(self):
        return f'{self.document.title} - Chunk {self.chunk_index}'
